﻿use crate::vm::InstructionType;
use  crate::parser::TreeNode;


 
/// Take the top-level program TreeNode and compile it to instructions.
#[must_use]
pub fn compile(ast: TreeNode) -> Vec<InstructionType> {
    let mut cg = Codegen::default();
    cg.compile(ast);
    cg.code
}

/// The Generator traverses the parsed source code and generates
/// `code` in the process.
#[derive(Default)]
struct Codegen {
    code: Vec<InstructionType>,
}

impl Codegen {
    #[allow(clippy::unused_self)]
    fn global(&self, v: &str) -> usize {
        match v {
            v if v.len() == 1 => v.chars().next().unwrap() as usize - 97,
            _ => panic!("{v} isn't a variable we can compile right now"),
        }
    }

    fn here(&self) -> usize {
        self.code.len()
    }

    fn hole(&mut self) -> usize {
        let p = self.here();
        self.code.push(InstructionType::Address(0));
        p
    }

    fn fix(&mut self, hole: usize, target: usize) {
        self.code[hole] = InstructionType::Address(target);
    }

    fn compile(&mut self, n: TreeNode) {
        match n {
            TreeNode::Add(a, b) => {
                self.compile(*a);
                self.compile(*b);
                self.code.push(InstructionType::Add);
            }
            TreeNode::Sub(a, b) => {
                self.compile(*a);
                self.compile(*b);
                self.code.push(InstructionType::Sub);
            }
            TreeNode::If1(test, then) => {
                self.compile(*test);
                self.code.push(InstructionType::Jz);
                let jz = self.hole();

                self.compile(*then);
                self.fix(jz, self.here());
            }
            TreeNode::If2(test, then, else_) => {
                self.compile(*test);
                self.code.push(InstructionType::Jz);
                let jz = self.hole();

                self.compile(*then);
                self.code.push(InstructionType::Jmp);
                let jmp = self.hole();

                self.fix(jz, self.here());
                self.compile(*else_);

                self.fix(jmp, self.here());
            }
            TreeNode::While(test, body) => {
                let l_restart = self.here();

                self.compile(*test);

                self.code.push(InstructionType::Jz);
                let jz = self.hole();

                self.compile(*body);
                self.code.push(InstructionType::Jmp);
                let jmp = self.hole();

                self.fix(jmp, l_restart);
                self.fix(jz, self.here());
            }
            TreeNode::Do(body, test) => {
                let l_restart = self.here();

                self.compile(*body);
                self.compile(*test);

                self.code.push(InstructionType::Jnz);
                let jnz = self.hole();
                self.fix(jnz, l_restart);
            }
            TreeNode::Prog(body) => {
                self.compile(*body);
                self.code.push(InstructionType::Halt);
            }
            TreeNode::Expr(body) => {
                self.compile(*body);
                self.code.push(InstructionType::Pop);
            }
            TreeNode::Set(var, expr) => {
                self.compile(*expr);
                self.code.push(InstructionType::Store);
                let TreeNode::Var(v) = *var else {
                    panic!("We expected a Var, not {:?}", *var);
                };
                self.code.push(InstructionType::Address(self.global(&v)));
            }
            TreeNode::Cst(val) => {
                self.code.push(InstructionType::Push);
                self.code.push(InstructionType::Integer(val));
            }
            TreeNode::Var(v) => {
                self.code.push(InstructionType::Fetch);
                self.code.push(InstructionType::Address(self.global(&v)));
            }
            TreeNode::Lt(a, b) => {
                self.compile(*a);
                self.compile(*b);
                self.code.push(InstructionType::Lt);
            }
            TreeNode::Seq(a, b) => {
                self.compile(*a);
                self.compile(*b);
            }
            TreeNode::Empty => {}
        }
    }
}
